﻿//#define UPDATE_HOSTS_FROM_SCHEMA_VERSIONS
#define UPDATE_HOSTS_FROM_DESCRIBE_RING
using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using Apache.Cassandra;

namespace Cassandraemon.Connection
{
	public class RingNodeEnumerator
		: ManualNodeEnumerator
	{
		public RingNodeEnumerator(ConnectionPool pool)
			: base(pool)
		{
			if (pool.Keyspace == "system") // There is no ring for the keyspace: system
			{
				this._DisableUpdateNodes = true;
			}
		}

		bool _DisableUpdateNodes;
		static readonly Random _Random = new Random();

		#region ICassandraNodeEnumerator

		public override void UpdateNodes()
		{
			if (this._DisableUpdateNodes)
			{
				return;
			}

			// Update TokenMap

			var pool = this.Pool;
			var client = pool.Borrow();
#if UPDATE_HOSTS_FROM_SCHEMA_VERSIONS
            Dictionary<string, List<string>> schemaVersions;
			try
			{
                schemaVersions = client.describe_schema_versions();
                pool.Return(client);
			}
            catch (Exception err)
            {
                CassandraContext.TraceError("describe_ring: {0}",
                    err);
                pool.ErrorReturn(client, err);
                throw;
            }
            var rghost = schemaVersions
                .SelectMany(sv => sv.Value)
                .Select(h => h.Trim())
                .Distinct()
                .ToList();
#elif UPDATE_HOSTS_FROM_DESCRIBE_RING
			List<TokenRange> rgring;
			try
            {
				rgring = client.describe_ring(pool.Keyspace);
                pool.Return(client);
            }
            catch (Exception err)
            {
                CassandraContext.TraceError("describe_ring: {0}",
                    err);
                if (err is InvalidRequestException)
                {
                    // InvalidRequestException could mean
                    // there is no ring for the keyspace.

					pool.Return(client);
                    this._DisableUpdateNodes = true;
                    return;
                }
				pool.ErrorReturn(client, err);
                throw;
            }

            var rghost = rgring
                .SelectMany(tr => tr.Endpoints)
                .Distinct()
                .ToList();
#else
			string tokenMapString;
			try
			{
				tokenMapString = client.get_string_property("token map");
				pool.Return(client);
			}
			catch (Exception err)
			{
				CassandraContext.TraceError("get_string_property(\"token map\") {0}",
					err);
				pool.ErrorReturn(client, err);
				throw;
			}

			// token map is a string like this:
			// {"0":"192.168.123.146","85070591730234615865843651857942052864":"192.168.123.142"}

			var rghost = tokenMapString
				.Replace("{", "")
				.Replace("}", "")
				.Replace("\"", "")
				.Split(',')
				.Select(tokenPair =>
				{
					var token = tokenPair.Split(':');
					return token[1].Trim();
				})
				.ToList();
#endif

			// There are two cases where token map should not update the hosts to use.
            // One is where the cluster uses internal network addresses,
            // which cannot reach from the current client.
            // The other is where storage.conf is not modified properly,
            // and therefore the ring may be something like:
			// {"58624523936365330278771871070307322155":"127.0.0.1"}
            // Check if any hosts have common IPAddress with the current settings,
            // and ignore if not.

            var rgadr = new HashSet<IPAddress>();
            foreach (var host in this.Nodes)
            {
                try
                {
                    var result = Dns.GetHostAddresses(host);
                    rgadr.UnionWith(result);
                }
                catch (SocketException err)
                {
                    // The host name is unresolvable.
                    // Not much we can do here,
                    // but probably safe not to update token map.

                    CassandraContext.TraceWarning(
                        "{0}: {1}",
                        host,
                        err);
                    this._DisableUpdateNodes = true;
                    return;
                }
            }

            rgadr.RemoveWhere(ipa => IPAddress.IsLoopback(ipa));
            if (rgadr.Count > 0)
            {
                var rgadrOfHost = rghost
                    .SelectMany(host => Dns.GetHostAddresses(host));
                rgadr.IntersectWith(rgadrOfHost);
                if (rgadr.Count == 0)
                {
                    CassandraContext.TraceWarning(
						"The ring has no common address with current settings, ignored"
#if UPDATE_HOSTS_FROM_SCHEMA_VERSIONS
#elif UPDATE_HOSTS_FROM_DESCRIBE_RING
#else
                        + ": {0}", tokenMapString
#endif
						);
					this._DisableUpdateNodes = true;
                    return;
                }
            }

			this.SetNodes(rghost);

            // Rotate by random count.
            // Without this, multiple programs that use single context gets the same connection.

            var cHost = this.Nodes.Count;
            if (cHost > 1)
            {
                for (var rnd = _Random.Next(cHost); rnd > 0; rnd--)
                {
                    this.RoundRobin();
                }
            }
		}

		#endregion
	}
}
